import datetime
import math
import os
import random

modify = 1
typeNameList = ["business","sport","auto"]

defualtRate = [0,0,0]    #每个类别未出现词的默认频率
typeNum = [0,0,0]        #各分类文章总数
typeRate = []            #各分类文章频率 P(yi)
articleNum = 0           #文章总数
typeId = 0

sourceDir = "F:\\data"
testFileList = []

count = 0
totalWords = 0
wordListValue = [0,0,0]  #各分类词汇总数
wordList = [{},{},{}]    #各分类词汇分别的数量
wordRate = [{},{},{}]    #P(xi|yi)

starttime = datetime.datetime.now()
print ("naive bayes algorithem start……")

#1、逐个文件训练，随机选一部分备测试
for fileName in os.listdir(sourceDir):
    
    count += 1
    if count % 1000 == 0 :
        print (str(count) + " files processed")
        
    #用于验证测试
    if random.randint(0,5) == 0:
        testFileList.append(fileName)
        continue

    #训练数据
    articleNum += 1
        
    #读取文章总数和类别
    for i in range(len(typeNameList)):
        if fileName.find(typeNameList[i]) > 0:
            typeNum[i] += 1
            typeId = i
            break
    
    page = open(sourceDir + "\\" + fileName,"r",encoding="UTF-8")

    #读取单词出现次数
    for lines in page.readlines():
        for word in lines.replace('\n', ' ').split(' '):
            word = word.strip()
            if len(word) < 1: continue

            totalWords += 1
            wordListValue[typeId] += 1

            if word in wordList[typeId]:
                wordList[typeId][word] += 1
            else:
                wordList[typeId][word] = 1

    page.close()

print (str(count) + " files processed\n")

#2、输出训练结果
for i in range(len(typeNameList)):
    print(typeNameList[i] + ":" + str(typeNum[i]))

for i in range(len(wordList)):
    print("wordList[" + str(i) + "]:" + str(len(wordList[i])))

print("articleNum:" + str(articleNum))
print("totalWords:" + str(totalWords))

#输出类别名
datamodel = open("data.model","w",encoding="UTF-8")

for name in typeNameList:
    datamodel.write(name + " ")
datamodel.write("\n")

#输出类别概率
for i in range(len(typeNameList)):
    rate = (typeNum[i]+0.0)/articleNum
    typeRate.append(rate)
    datamodel.write(str(rate) + " ")
datamodel.write("\n")

#输出字典及概率
wordSet = set({})
for i in range(len(wordList)):
    wordSet |= set(wordList[i].keys())

vocabulary = len(wordSet)

print("vocabulary:" + str(vocabulary))

for i in range(len(wordList)):
    defualtRate[i] = modify/(wordListValue[i] + vocabulary * modify)
    datamodel.write(str(defualtRate[i]) + "\n")

    for k,v in wordList[i].items():
        datamodel.write(k + " ")
        r = (v + modify)/(wordListValue[i] + vocabulary * modify)
        wordRate[i][k] = r
        datamodel.write(str(r) + " ")
    datamodel.write("\n")

datamodel.close()

for i in range(len(wordRate)):
    print("sum(wordRate["+str(i)+"]):" + str(sum(wordRate[i].values())))
    
print ("train end!\n")


#3、验证模型
accuracyNum = 0        #准确数
precisionNum= [0,0,0]  #精确数
recallNum = [0,0,0]    #召回数

accuracyTotal = 0
precisionTotal = [0,0,0]
recallTotal = [0,0,0]

accuracyTotal = len(testFileList)

predictResult = open("data.result", "w", encoding="UTF-8")

#测试每一个文件
for fileName in testFileList:
    uniqueWords = {}
    allWords = []
    maxPosible = 0
    maxType = 0

    for i in range(len(typeNameList)):
        if fileName.find(typeNameList[i]) > 0:
            realType = i
            break

    #将文件加载到allWords
    page = open(sourceDir + "\\" + fileName, "r", encoding="UTF-8")
    
    for lines in page.readlines():
        for word in lines.replace('\n', ' ').split(' '):
            if len(word.strip()) < 1: continue

            allWords.append(word)

    page.close()

    #根据模型进行分类
    for i in range(len(wordRate)):
        curPosible = math.log(typeRate[i])

        for word in allWords:
            curPosible = curPosible + math.log(wordRate[i].get(word,defualtRate[i]))

        if i == 0 or maxPosible < curPosible:
            maxPosible = curPosible
            maxType = i


    if realType == maxType:
        accuracyNum += 1
        recallNum[realType] += 1
        precisionNum[maxType] += 1

    recallTotal[realType] += 1
    precisionTotal[maxType] += 1
    
    #for future use
    predictResult.write(str(realType) + " " + str(maxType) + "\n")

predictResult.close()

#输出结果
print("accurate:" + str((accuracyNum+0.0)/accuracyTotal))

for i in range(3):
    print(str(i) + "\'s precision, recall: " + str((precisionNum[i]+0.0)/precisionTotal[i]) + "," + str((recallNum[i]+0.0)/recallTotal[i]))
    
endtime = datetime.datetime.now()
s = (endtime - starttime).seconds
print ("\ntime cost: "+str(s) + "s")
